import os
import torch
import utils
from torch.autograd import Variable
import torch.nn.functional as F
import math
import numpy as np
import tool as tl
import model.net_2 as net2

path = 'epoch_params/60epoch/'

def get_ie(model, checkpoint, params, val_dataloader, device):
    utils.load_checkpoint(checkpoint, model)
    model.eval()
    # summary for current eval loop
    summ = 0
    for data_batch, labels_batch in val_dataloader:
        data_batch, labels_batch = Variable(data_batch), Variable(labels_batch)
        # compute model output
        with torch.no_grad():
            output_teacher_batch = model(data_batch)
        out = F.softmax(output_teacher_batch)
        for idx, label in enumerate(labels_batch):
            # print(out[idx])
            tmp = out[idx]
            for p in out[idx]:
                summ += p*math.log(p)
    return summ

def label_ave_num(path, model, params, val_dataloader, device):
    circle = os.listdir(path)
    label_ie = []
    for ci in circle:
        print(ci)
        circle_path = path + ci + '/'
        ie_0_path = circle_path + '-1' + '/' + 'best.pth.tar'
        ie_init = get_ie(model, ie_0_path, params, val_dataloader, device)
        ie_list = []
        for i in range(60):
            ie_1_path = circle_path + str(i) + '/' + 'best.pth.tar'
            ie = get_ie(model, ie_1_path, params, val_dataloader, device)
            ie_list.append(ie - ie_init)
        label_ie.append(ie_list)
    label_ie_ave = list(tl.get_mean_arr(np.array(label_ie)))
    return label_ie_ave


def ie_picnum(path, model, params, val_dataloader, device, num):
    label_names = os.listdir(path)
    label = []
    curvedatas = []
    name = label_names[num]
    label.append(name)
    print('ie: {} beginning'.format(name))
    label_path = path + name + '/'

    label_ie_ave = label_ave_num(label_path, model, params, val_dataloader, device)
    curvedatas.append(label_ie_ave)

    print('ie: {} finish'.format(name))

    return curvedatas, label


# torch.cuda.set_device(1)
# device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
#
# json_path = 'experiments/base_cnn/params.json'
# params = params = utils.Params(json_path)
# params.cuda = torch.cuda.is_available()
# print('the gpu is {}'.format(params.cuda))
# print('the current gpu is {}'.format(torch.cuda.current_device()))
# model = nn2.Net(params)